
/**
 * @file MultigridSolver.h
 * @brief The abstract class about multigrid solver
 * @author XDDDD
 * @version 
 * @date 2021-06-10
 */
#ifndef __PAISLEYPARK__MULTIGRIDSOLVER_H__
#define __PAISLEYPARK__MULTIGRIDSOLVER_H__

#include <iostream>
#include <vector>
#include "RestrictionOperatorFactory.h"
#include "InterpolationOperatorFactory.h"
#include "MultigridVector.h"
#include "VectorFunction.h"
#include <lapacke.h>

template <size_t DIM>
class MultigridSolver {
public:
	/**
	 * @brief Construction
	 */
	MultigridSolver();
	/**
	 * @brief Construction with restriction operator and interpolation operator.
	 *
	 * @param _RestrictionId  The Id of the restriction operator.
	 * @param _InerpolationId The Id of the interpoaltion operator.
	 */
	MultigridSolver(std::string _RestrictionId, std::string _InerpolationId);
	/**
	 * @brief Destruction
	 */
	~MultigridSolver(){};
	/**
	 * @brief Set the grid.
	 *
	 * @param _v The given initial guess.
	 * @param _N The order of the finest grid.
	 * @param _f The equation function.
	 * @param _g The boundary function.
	 */
	void set(const std::vector<double>& _v, const size_t _N, VectorFunction<DIM>& _f, VectorFunction<DIM>& _g);	
	/**
	 * @brief Set the grid.
	 *
	 * @param _v The given initial guess.
	 * @param _N The order of the finest grid.
	 * @param _f The equation funciton.
	 */
	void set(const std::vector<double>& _v, const size_t _N, VectorFunction<DIM>& _f);
	/**
	 * @brief Get the order of the finest grid.
	 *
	 * @return __N
	 */
	const size_t get_N() const;
	/**
	 * @brief Get the value of one knot of __U. 
	 *
	 * @param _coord The coord of the knot.
	 * @param _N The order of the grid.
	 *
	 * @return The value of the knot
	 */
	const double get_U(const std::vector<size_t>& _coord, const size_t _N) const;
	/**
	 * @brief Get the value of one knot of __F.
	 *
	 * @param _coord The coord of the knot.
	 * @param _N The order of the grid.
	 *
	 * @return The value of the knot
	 */
	const double get_F(const std::vector<size_t>& _coord, const size_t _N) const;
	/**
	 * @brief Get the value of one knot of __U
	 *
	 * @param _idx The index of the knot.
	 * @param _N The order of the grid.
	 *
	 * @return The value of the knot.
	 */
	const double get_U(const size_t _idx, const size_t _N) const;
	/**
	 * @brief Get the value of one knot of __F
	 *
	 * @param _idx The index of the knot.
	 * @param _N The order of the grid.
	 *
	 * @return The value of the knot.
	 */
	const double get_F(const size_t _idx, const size_t _N) const;
	/**
	 * @brief Smooth on the given grid.
	 *
	 * @param _M The order of the grid.
	 */
	void smooth(const size_t _M);
	/**
	 * @brief Get the precise solution on the coarest grid.
	 *
	 * @param _M The order of the given grid.
	 */
	void solve_coarest(const size_t _M);
	/**
	 * @brief Set f,
	 *
	 * @param _f The equation function
	 * @param _g The boundary function
	 * @param _M The order of the grid.
	 */
	void set_f(VectorFunction<DIM>& _f, VectorFunction<DIM>& _g, const size_t _M);
	/**
	 * @brief Set f
	 *
	 * @param _f The equation function.
	 * @param _M The order of the grid
	 */
	void set_f(VectorFunction<DIM>& _f, const size_t _M);
	/**
	 * @brief Set the error back to __F and restrict it to coareser grid.
	 *
	 * @param _M The given order of present grid.
	 */
	void set_e(const size_t _M);
	/**
	 * @brief Set the error of v back to __U and interpolate it to finer grid.
	 *
	 * @param _M The given order of present grid.
	 */
	void putback_e(const size_t _M);
	/**
	 * @brief Get the error of the solution with infinite norm
	 *
	 * @param _g The boundary function
	 *
 	 * @return  The error
	 */
	const double get_E(VectorFunction<DIM>& _g) const;
	/**
	 * @brief Solve the poisson equation.
	 *
	 * @param _f The equation function.
	 * @param _g The boundary function.
	 * @param _v The initial guess.
	 * @param _N The order of the finest grid.
	 * @param _M The order of the coarest grid.
	 * @param itr_num The number that we perform iteration.
	 * @param nu1 The number of the frist smooth iteration.
	 * @param nu2 The number of the second smooth iteration.
	 * @param _epsilon The tolerant error.
	 */
	virtual void Solve(VectorFunction<DIM>& _f, VectorFunction<DIM>& _g, std::vector<double>& _v,
					   const size_t _N, const size_t _M, const size_t itr_num,  const size_t nu1, const size_t nu2, const double _epsilon) = 0;
protected:

	/**
	 * @brief The interpolation operator
	 */
	InterpolationOperator<DIM>* __I_2h_2_h;

	/**
	 * @brief The restriction operator
	 */
	RestrictionOperator<DIM>* __I_h_2_2h;
	

	/**
	 * @brief n=2^N, where n is the number of grids in each dimension.
	 */
	size_t __N = 0;			
	/**
	 * @brief Store the initial guess and the solution will replace it after solving the PDE.
	 */
	MultigridVector<DIM> __U;

	/**
	 * @brief Store the right side of the equation and residual equation.
	 */
	MultigridVector<DIM> __F;

};

template <size_t DIM>
MultigridSolver<DIM>::MultigridSolver() {
	__N = 0;
	__F.clear();
	__U.clear();
	__I_2h_2_h = InterpolationOperatorFactory<DIM>::Instance()->CreateOperator("Linear");
	__I_h_2_2h = RestrictionOperatorFactory<DIM>::Instance()->CreateOperator("Fullweighting");
};

template <size_t DIM>
MultigridSolver<DIM>::MultigridSolver(std::string _RestrictionId, std::string _InterpolationId) {
	__N = 0;
	__F.clear();
	__U.clear();
	__I_2h_2_h = InterpolationOperatorFactory<DIM>::Instance()->CreateOperator(_InterpolationId);
	__I_h_2_2h = RestrictionOperatorFactory<DIM>::Instance()->CreateOperator(_RestrictionId);
};

template <size_t DIM>
void MultigridSolver<DIM>::set(const std::vector<double>& _v, const size_t _N, VectorFunction<DIM>& _f, VectorFunction<DIM>& _g) {
	size_t l = 1;
	for(size_t i = 0; i != DIM; i++) {
		l = l*(1 << _N) + l;
	}
	if(_v.size() != l) {
		std::cout << "The length does not meet with the initial guess!" << std::endl;
		std::exit(-1);
	}
	__N = _N;
	__U.set(_v, _N);
	__F.set(_v, _N);
	set_f(_f, _g, _N);
};

template <size_t DIM>
void MultigridSolver<DIM>::set(const std::vector<double>& _v, const size_t _N, VectorFunction<DIM>& _f) {
	size_t l = 1;
	for(size_t i = 0; i != DIM; i++) {
		l = l*(1 << _N) + l;
	}
	if(_v.size() != l) {
		std::cout << "The length does not meet with the initial guess!" << std::endl;
		std::exit(-1);
	}
	__N = _N;
	__U.set(_v, _N);
	__F.set(_v, _N);
	set_f(_f, _N);
};

template <size_t DIM>
const size_t MultigridSolver<DIM>::get_N() const {
	return __N;
};

template <size_t DIM>
const double MultigridSolver<DIM>::get_U(const std::vector<size_t>& _coord, const size_t _N) const {
	return __U(_coord, _N);
};

template <size_t DIM>
const double MultigridSolver<DIM>::get_F(const std::vector<size_t>& _coord, const size_t _N) const {
	return __F(_coord, _N);
};

template <size_t DIM>
const double MultigridSolver<DIM>::get_U(const size_t _idx, const size_t _N) const {
	return __U(_idx, _N);
};

template <size_t DIM>
const double MultigridSolver<DIM>::get_F(const size_t _idx, const size_t _N) const {
	return __F(_idx, _N);
};

template <size_t DIM>
void MultigridSolver<DIM>::smooth(const size_t _M) {
	double w = 2.0/3;
	if(DIM == 1) 
		w = 2.0/3;
	else if(DIM == 2)
		w = 4.0/5.0;
	size_t l = 1;
	size_t N = __U.get_N();
	size_t num_1d = (1 << _M) + 1;
	for(size_t i = 0; i != DIM; i++) {
		l = l*(1 << _M) + l;
	} 

	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	/*for(size_t i = 0; i != DIM; i++) {
		coord[i] = 1;
	}*/
	size_t pole = 0;
	size_t pre_index = 0;
	double coef = 1.0;
	coef = 0.5/DIM;
//	double v[l] = {0};
	std::vector<double> v(l, 0);
	bool is_stop;
	bool is_bound;
	for(size_t i = 0; i != 2; i++) {
		//Initialization.
		is_stop = false;
		pole = 0;
		l = 1;
		pre_index = 0;
		while(!is_stop) {
			pole = 0;
			is_bound = false;
			for(size_t j = 0; j != DIM; j++) {
				if(coord[j] == 0 || coord[j] == num_1d - 1)
					is_bound = true;
			}
			if(i == 1) {
				//Put the data back to __U
				__U(pre_index, _M) = v[pre_index];
				__U(pre_index, _M) += __F(pre_index, _M)*w*coef/(1 << _M + _M);
			}
			else {
				if(!is_bound) {
					v[pre_index] += __U(coord, _M)*(1.0 - w);
					for(size_t j = 0; j != DIM; j++) {
						coord[j] -= 1;
						v[pre_index] += w*coef*__U(coord, _M);
						coord[j] += 1;
						coord[j] += 1;
						v[pre_index] += w*coef*__U(coord, _M);
						coord[j] -= 1;
					}
				}
				else {
					v[pre_index] = __U(coord, _M);
				}
			}
			coord[pole]++;
			while(coord[pole] == num_1d) {
				coord[pole] = 0;
				pole++;
				if(pole == DIM)
					break;
				coord[pole]++;
			}
			if(pole == DIM)
				is_stop = true;
			else {
				pre_index++;
			}
		}
	}
};

template <size_t DIM>
void MultigridSolver<DIM>::solve_coarest(const size_t _M) {
	size_t l = 1;
	for(size_t i = 0; i != DIM; i++) {
		l *= (1 << _M) - 1;
	}
	
	size_t num_1d = (1 << _M) - 1;
	double A[l*l] = {0};
	double v[l] = {0};
	size_t d = 0;
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	for(size_t i = 0; i != DIM; i++) {
		coord[i] = 1;
	}
	size_t pre_index = 0;
	size_t pole = 0;
	bool is_bound;
	while(pole != DIM) {
		pole = 0;
		v[pre_index] += __F(coord, _M);
		for(size_t i = 0; i != DIM; i++) {
			if(coord[i] == 1) {
				coord[i]--;
				v[pre_index] += (1 << _M + _M)*__U(coord, _M);
				coord[i]++;
			}
			else if(coord[i] == num_1d) {
				coord[i]++;
				v[pre_index] += (1 << _M + _M)*__U(coord, _M);
				coord[i]--;
			}
		}
		coord[pole]++;
		pre_index++;
		while(coord[pole] == num_1d + 1) {
			coord[pole] = 1;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
	for(size_t i = 0; i != l; i++) {
		A[i*l + i] = 1.0*(1 << _M + _M)*2*DIM;
		for(size_t j = 0; j != DIM; j++) {
			if(j == 0)
				d = 1;
			else 
				d = d*num_1d;
			if(i >= d) 
				A[i*l + i - d] = -1.0*(1 << _M + _M);
			if(i + d < l) 
				A[i*l + i + d] = -1.0*(1 << _M + _M);
		}	
	}
	//solve the equations.
	int ipiv[l];
	LAPACKE_dgesv(LAPACK_COL_MAJOR, l, 1, A, l, ipiv, v, l);
	
	pole = 0;
	pre_index = 0;
	while(pole != DIM) {
		pole = 0;
		__U(coord, _M) = v[pre_index];
		coord[pole]++;
		pre_index++;
		while(coord[pole] == num_1d + 1) {
			coord[pole] = 1;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
};

template <size_t DIM>
void MultigridSolver<DIM>::set_f(VectorFunction<DIM>& _f, VectorFunction<DIM>& _g, const size_t _M) {
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	size_t pole = 0;
	size_t index = 0;
	bool is_bound;
	size_t num_1d = (1 << _M) + 1;
	std::vector<double> x;
	x.resize(DIM);
	__F.clean(_M);
	while(pole != DIM) {
		is_bound = false;
		for(size_t i = 0; i != DIM; i++) {
			if(coord[i] == 0 || coord[i] == num_1d - 1)
				is_bound = true;
		}
		if(is_bound) {
			for(size_t j = 0; j != DIM; j++) 
				x[j] = 1.0/(1 << _M)*coord[j];
			__U(coord, _M) = _g(x);
			__F(coord, _M) = 0;
		}
		else {
			for(size_t j = 0; j != DIM; j++) 
				x[j] = 1.0/(1 << _M)*coord[j];
			__F(coord, _M) = _f(x);
		}
		pole = 0;
		coord[pole]++;
		while(coord[pole] == num_1d) {
			coord[pole] = 0;
			pole++;
			if(pole == DIM) {
				break;
			}
			coord[pole]++;
		}
	}
};

template <size_t DIM>
void MultigridSolver<DIM>::set_f(VectorFunction<DIM>& _f, const size_t _M) {
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	size_t pole = 0;
	size_t index = 0;
	bool is_bound;
	size_t num_1d = (1 << _M) + 1;
	std::vector<double> x;
	x.resize(DIM);
	__F.clean(_M);
	while(pole != DIM) {
		is_bound = false;
		for(size_t i = 0; i != DIM; i++) {
			if(coord[i] == 0 || coord[i] == num_1d - 1)
				is_bound = true;
		}
		if(!is_bound) {
			for(size_t j = 0; j != DIM; j++) 
				x[j] = 1.0/(1 << _M)*coord[j];
			__F(coord, _M) = _f(x);
		}
		pole = 0;
		coord[pole]++;
		while(coord[pole] == num_1d) {
			coord[pole] = 0;
			pole++;
			if(pole == DIM) {
				break;
			}
			coord[pole]++;
		}
	}
};

template <size_t DIM>
void MultigridSolver<DIM>::set_e(const size_t _M) {
	size_t l = 1;
	for(size_t i = 0; i != DIM; i++) {
		l *= (1 << _M) + 1;
	}
	size_t num_1d = (1 << _M) + 1;
	
	std::vector<double> F(l, 0);
	//double F[l] = {0};
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	size_t pole = 0;
	double mid;

	for(size_t i = 0; i != l; i++) {
		F[i] = __F(i, _M);
	}
	for(size_t i = 0; i != DIM; i++) {
		coord[i] = 1;
	}
	while(pole != DIM) {
		pole = 0;
		__F(coord, _M) -= 2.0*DIM*(1 << _M + _M)*__U(coord, _M);
		for(size_t i = 0; i != DIM; i++) {
			coord[i]--;
			mid = 1.0*(1 << _M + _M)*__U(coord, _M);
			coord[i]++;
			__F(coord, _M) += mid;
			coord[i]++;
			mid = 1.0*(1 << _M + _M)*__U(coord, _M);
			coord[i]--;
			__F(coord, _M) += mid;
		}
		coord[pole]++;
		while(coord[pole] == num_1d - 1) {
			coord[pole] = 1;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
	(*__I_h_2_2h)(__F, _M);
	for(size_t i = 0; i != l; i++) {
		__F(i, _M) = F[i];
	}
};

template <size_t DIM>
void MultigridSolver<DIM>::putback_e(const size_t _M) {
	size_t l = 1;
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	for(size_t i = 0; i != DIM; i++) {
		l *= (1 << _M + 1) + 1;
	}
	
	std::vector<double> v(l, 0);
	//double v[l] = {0};
	for(size_t i = 0; i != l; i++) {
		v[i] = __U(i, _M + 1);
	}
	(*__I_2h_2_h)(__U, _M);

	size_t pole = 0;
	bool is_jump = false;
	size_t index = 0;
	while(pole != DIM) {
		is_jump = false;
		pole = 0;
		for(size_t i = 0; i != DIM; i++) {
			if(coord[i] == 0 || coord[i] == (1 << _M + 1)) {
				is_jump = true;
				break;
			}
		}
		if(!is_jump)
			__U(index, _M + 1) += v[index];
		coord[pole]++;
		index++;
		while(coord[pole] == (1 << _M + 1) + 1) {
			coord[pole] = 0;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
};

template <size_t DIM>
const double MultigridSolver<DIM>::get_E(VectorFunction<DIM>& _g) const{
	double E = 0;
	std::vector<size_t> coord(DIM, 1);
	std::vector<double> x(DIM, 0);
	size_t pole = 0;
	size_t N = __N;
	bool is_bound = false;
	while(pole != DIM) {
		is_bound = false;
		pole = 0;
		for(size_t i = 0; i != DIM; i++) {
			x[i] = 1.0*coord[i]/(1<<N);
		}
		if(E < fabs(__U(coord, N) - _g(x))) {
			E = fabs(__U(coord, N) - _g(x));			
		}

		coord[pole]++;
		while(coord[pole] == (1 << N)) {
			coord[pole] = 1;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
	return E;
};

#else
//Do nothing.
#endif
